package com.liuyu.common.framework.lang.util.ratelimiter;

import java.time.LocalDateTime;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * 漏桶算法
 * <p>
 * 可以理解为注水漏水的过程，往漏桶中以任意速率流入水，以固定的速率流出水。
 * 当水超过桶的容量时，会被溢出，也就是被丢弃。因为桶容量是不变的，保证了整体的速率。
 * 流入的水滴，可以看作是访问系统的请求，这个流入速率是不确定的；
 * 桶的容量一般用来表示系统所能处理的请求数；如果桶的容量满了，也就达到限流的阀值，会丢弃水滴（即：拒绝请求）；
 * 流出的水滴，是恒定速率的，用来表示服务按照固定的速率处理请求。
 *
 * @author: liuyu
 * @date: 2023/5/22 8:51
 * @version: 1
 */
public class LeakyBucketRateLimiter {

    private static final int INT = 20;
    private int qps = 2;

    private LinkedBlockingQueue<Character> waterBucket;


    /**
     * 表示桶是否被初始化
     */
    private volatile AtomicBoolean initialized = new AtomicBoolean(false);

    /**
     * 关联一个漏桶（一般情况为自身）
     */
    private static volatile LeakyBucketRateLimiter leakyBucketRateLimiter;


    public static LeakyBucketRateLimiter getLeakyBucket() {
        return leakyBucketRateLimiter;
    }

    private LeakyBucketRateLimiter(int capacity, int qps) {
        // 漏桶只能被初始化一次
        if (initialized.compareAndSet(false, true)) {
            this.qps = qps;
            waterBucket = new LinkedBlockingQueue<>(capacity);
            // 初始化消费者
            initConsumer();
        }
    }

    /**
     * 漏桶中的水以固定速率流出
     */
    private void initConsumer() {
        new Thread(() -> {
            while (true) {
                try {
                    TimeUnit.MILLISECONDS.sleep(1000 / qps);
                } catch (InterruptedException e) {
                    // log for exception
                    System.out.println("Exception occur! " + e);
                }
                // 以固定速率消费
                waterBucket.poll();
            }
        }).start();
    }


    /**
     * 是否将请求加入到漏桶，能存入则代表漏桶没满，允许请求通过，返回true；否则返回false
     *
     * @return
     */
    public boolean tryAcquire() {
        return waterBucket.offer('S');
    }


    /**
     * 创建漏铜
     *
     * @param capacity 漏桶容量
     * @param qps      漏桶支持的QPS
     * @return com.liuyu.common.framework.lang.util.ratelimiter.LeakyBucketRateLimiter
     * @throws
     * @author liuyu
     * @date 2023/5/22 8:53
     */
    public static LeakyBucketRateLimiter create(int capacity, int qps) {
        leakyBucketRateLimiter = new LeakyBucketRateLimiter(capacity, qps);
        return leakyBucketRateLimiter;
    }


    /**
     * 漏桶容量2，每秒流出2滴水；
     * 最初漏桶中没有水，所以第一秒立刻打进去两滴水；又由于漏桶每秒流出2滴水，所以在程序开始跑时，第一秒必回流出一滴水（极限情况2滴）
     * 所以第一秒会进入3滴水，第一秒之后每秒稳定流入2滴水，即QPS为2；
     *
     * @param args
     */
    public static void main(String[] args) {
        LeakyBucketRateLimiter leakyBucket = LeakyBucketRateLimiter.create(2, 2);

        for (int i = 0; i < INT; i++) {
            LocalDateTime now = LocalDateTime.now();
            if (leakyBucket.tryAcquire()) {
                System.out.println(now + " pass the rate limiting");
            } else {
                System.out.println(now + " was limited");
            }

            try {
                TimeUnit.MILLISECONDS.sleep(200);
            } catch (InterruptedException e) {
                // 日志记录
            }

        }

        System.out.println("------------");
        // 再次获取漏桶
        LeakyBucketRateLimiter leakyBucket2 = LeakyBucketRateLimiter.getLeakyBucket();

        // 验证漏桶只会有一个
        System.out.println("leakyBucket only one ? " + (leakyBucket == leakyBucket2));

        for (int i = 0; i < 10; i++) {
            LocalDateTime now = LocalDateTime.now();
            if (leakyBucket2.tryAcquire()) {
                System.out.println(now + " pass the rate limiting");
            } else {
                System.out.println(now + " was limited");
            }

            try {
                TimeUnit.MILLISECONDS.sleep(200);
            } catch (InterruptedException e) {
                // 日志记录
            }

        }
    }
}
